from sklearn.model_selection import train_test_split
import numpy as np
import os

if __name__ == "__main__":
    path = '.'
    ad_array, pwm_array = [], []
    files = os.listdir(path)

    # 加载数据
    for file in files:
        if file.split('.')[-1] == 'npy':
            if file[0:9] == 'origin_ad':
                ad_data = np.load(file)
                if len(ad_array) == 0:
                    ad_array = ad_data
                else:
                    ad_array = np.concatenate((ad_array, ad_data))
            elif file[0:10] == 'origin_pwm':
                pwm_data = np.load(file)
                if len(ad_array) == 0:
                    pwm_array = pwm_data
                else:
                    pwm_array = np.concatenate((pwm_array, pwm_data))

    x_ad_data = ad_array
    y_pwm_data = pwm_array

    out_len = int(len(x_ad_data)/100)*100
    ad_npy = x_ad_data[0:out_len]
    label_pwm = y_pwm_data[0:out_len]
    res_ad = x_ad_data[out_len:]
    res_pwm = y_pwm_data[out_len:]

    # 划分训练集和测试集
    train_data, test_data, train_label, test_label = train_test_split(ad_npy, 
                                                                      label_pwm,
                                                                      train_size=0.8,
                                                                      random_state=0)    
    train_data = np.concatenate((train_data, res_ad))
    train_label = np.concatenate((train_label, res_pwm))

    # 保存训练集和测试集
    np.save('./ad_train_data.npy', train_data)
    np.save('./ad_test_data.npy', test_data)
    np.save('./pwm_train_label.npy', train_label)
    np.save('./pwm_test_label.npy', test_label)

    print('ad train data shape:%f~%f' % (min(train_data.flatten()), max(train_data.flatten())))
    print('pwm train data shape:%f~%f' % (min(train_label), max(train_label)))
    print('ad test data shape:%f~%f' % (min(test_data.flatten()), max(test_data.flatten())))
    print('pwm test data shape:%f~%f' % (min(test_label), max(test_label)))
    print('Generate ad_train_dat.npy,ad_test_dat.npy,pwm_train_label.npy,pwm_test_label.npy')